from agentDesign.L1.perception_module import ReadMap
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
import math

# import plotly.graph_objects as go
import numpy as np

map_file = "./L1/map.csv"
building_file = "./L1/building_code.csv"
rider_num = 18
day = 8


def draw_heap_map():
    read_map = ReadMap(map_file, building_file)
    map_matrix = read_map.map_matrix
    heat_matrix = np.zeros((map_matrix.shape[0], map_matrix.shape[1]))

    for i in range(rider_num):
        # deliver_data = database.search_data(
        #     # table_name=f"deliver_{i}_record_regular",
        #     # table_name=f"deliver_{i}_record_imitation",
        #     # table_name=f"deliver_{i}_record_qlearning",
        #     table_name=f"deliver_{i}_record_qlmemory",
        #     # table_name=f"deliver_{i}_record_mix",
        #     column="*"
        # )
        # filename = f"deliver_{i}_record_regular.csv"
        # filename = f"deliver_{i}_record_imitation.csv"
        filename = f"./llm_agent_log/deliver_{i}_record_llm.csv"
        data_frame = pd.read_csv(filename)
        for index, data in data_frame.iterrows():
            if data.iloc[5] == 'a':
                continue
            if data.iloc[11] <= day / 2:
                continue
            # cnt += 1
            # print(cnt)
            heat_matrix[data.iloc[1]][data.iloc[2]] += 1

    divide_rate = 5
    small_matrix = np.zeros(
        (math.ceil(map_matrix.shape[0] / divide_rate), math.ceil(map_matrix.shape[1] / divide_rate)))
    for i in range(map_matrix.shape[0]):
        for j in range(map_matrix.shape[1]):
            small_matrix[int(i / divide_rate)][int(j / divide_rate)] += heat_matrix[i][j]

    merchants = ["饭店", "百货大楼", "咖啡店", "药房", "书店", "面包店", "水果店"]
    merchants_x = []
    merchants_y = []
    for merchant in merchants:
        merchant_doors = read_map.get_door(merchant)
        for door in merchant_doors:
            merchants_x.append(door[0][0])
            merchants_y.append(door[0][1])
    max_x = max(merchants_x)
    max_y = max(merchants_y)

    small_matrix = small_matrix[2: 34, :32]

    heatmap = sns.heatmap(small_matrix, cmap="coolwarm")
    for i in range(len(merchants_x)):
        y_mark, x_mark = int(merchants_x[i] / divide_rate), int(merchants_y[i] / divide_rate)  # 注意坐标系中x和y的位置，这里与数组索引相反
        heatmap.text(x_mark + 0.5, y_mark + 0.5, '\u2605', color='DarkOrange', ha='center', va='center', fontsize=10)

    # plt.title('Heatmap of Courier(Regular) Trajectories Activity Areas',  font={'family': 'Arial', 'size': 12})
    # plt.title('Heatmap of Courier(Imitation) Trajectories Activity Areas',  font={'family': 'Arial', 'size': 12})
    # plt.title('Heatmap of Courier(RL) Trajectories Activity Areas',  font={'family': 'Arial', 'size': 12})
    # plt.title('Heatmap of Courier(RL+Memory) Trajectories Activity Areas', font={'family': 'Arial', 'size': 12})
    plt.title('Heatmap of Courier(llm) Trajectories Activity Areas2', font={'family': 'Arial', 'size': 12})
    # plt.title('Heatmap of Courier(Mix) Trajectories Activity Areas',  font={'family': 'Arial', 'size': 12})
    # plt.savefig("../PlotResult/complete_plot/heatmap_regular.png", dpi=300)
    # plt.savefig("../PlotResult/complete_plot/heatmap_imitation.png", dpi=300)
    # plt.savefig("../PlotResult/complete_plot/heatmap_RL.png", dpi=300)
    # plt.savefig("../PlotResult/complete_plot/heatmap_RL_memory.png", dpi=300)
    # plt.savefig("../PlotResult/complete_plot/heatmap_mix.png", dpi=300)
    plt.savefig("./Storage/llm_result_image/heatmap_llm_2.png", dpi=300)
    # plt.show()

    # fig = go.Figure(data=go.Heatmap(
    #     z=map_matrix,
    #     hoverongaps=False))
    # fig.show()


if __name__ == '__main__':
    draw_heap_map()
